"""
PostgreSQL增强恢复模块
包含完整的数据重组、解密、解压处理流程
"""

import os
import json
import logging
from modules import CompressionModule, EncryptionModule

LOGGER = logging.getLogger(__name__)

class RestoreManager:
    def __init__(self, config):
        """
        初始化恢复管理器
        :param config: 配置字典，包含：
            - backup_dir: 备份存储目录
            - encryption_key: 加密密钥（需与备份时一致）
        """
        self.backup_dir = config['backup_dir']
        self.encryption = EncryptionModule(config['encryption_key'])
        
        LOGGER.debug("RestoreManager initialized")

    def _rebuild_file(self, manifest_path, target_path):
        """
        根据清单文件重建单个文件（内部方法）
        """
        try:
            # 1. 读取清单文件
            with open(manifest_path, 'r') as f:
                manifest = json.load(f)
            
            # 2. 初始化解压模块
            compression = CompressionModule(manifest['compression'])
            
            # 3. 重组数据块
            restored_data = bytearray()
            for block_hash in manifest['block_hashes']:
                # 构建块存储路径
                block_path = os.path.join(
                    self.backup_dir,
                    'blocks',
                    block_hash[:2],
                    block_hash[2:4],
                    block_hash
                )
                
                # 读取加密块
                with open(block_path, 'rb') as f:
                    encrypted_data = f.read()
                
                # 解密处理
                decrypted_data = self.encryption.decrypt(encrypted_data)
                
                # 解压处理
                decompressed_data = compression.decompress(decrypted_data)
                
                restored_data.extend(decompressed_data)
            
            # 4. 写入恢复文件
            os.makedirs(os.path.dirname(target_path), exist_ok=True)
            with open(target_path, 'wb') as f:
                f.write(restored_data)
            
            # 验证数据完整性
            actual_size = len(restored_data)
            expected_size = manifest['original_size']
            if actual_size != expected_size:
                raise ValueError(f"Size mismatch: expected {expected_size}, got {actual_size}")
            
            LOGGER.info(f"Successfully restored {target_path} | "
                       f"Original size: {expected_size} | "
                       f"Restored size: {actual_size}")
            return True
        except Exception as e:
            LOGGER.error(f"Failed to restore {target_path}: {str(e)}")
            return False

    def execute_restore(self, restore_mapping):
        """
        执行恢复操作
        :param restore_mapping: 恢复路径映射字典 {备份文件: 目标路径}
        :return: 成功恢复的文件数
        """
        success_count = 0
        
        # 遍历所有清单文件
        manifests_dir = os.path.join(self.backup_dir, 'manifests')
        for manifest_file in os.listdir(manifests_dir):
            manifest_path = os.path.join(manifests_dir, manifest_file)
            
            # 获取对应的目标路径
            with open(manifest_path, 'r') as f:
                manifest = json.load(f)
            original_path = manifest['original_path']
            target_path = restore_mapping.get(original_path)
            
            if not target_path:
                LOGGER.warning(f"No target mapping for {original_path}")
                continue
            
            if self._rebuild_file(manifest_path, target_path):
                success_count += 1
        
        LOGGER.info(f"Restore completed | Total: {len(restore_mapping)} | Success: {success_count}")
        return success_count